
# Copyright 2009 David Lee Miller
#
# This file is part of POSSM (Python Open Source Scanning Microscopy).
#
# POSSM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POSSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License         
# along with POSSM.  If not, see <http://www.gnu.org/licenses/>.
#                                                                           
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

from Scientific.IO import NetCDF
import Numeric as N
from controller import n,os,time,ws

class Class2NetCDF(object):
	''' A class for saving aquired data to standard NetCDF format. '''
	basedir = ws.BASE_DIRECTORY
	datadir = ws.DATA_DIRECTORY
	extension = ws.EXTENSION
	stm_prefix = ws.STM_PREFIX
	imagescans = ['image','imagewiv','imagewiz','imagewzv','multibias','continuous']
	ivscans = ['singleiv','imagewiv']
	izscans = ['singleiz','imagewiz']
	zvscans = ['singlezv','imagewzv']
	spectroscopy = ['imagewiv','imagewiz','imagewzv']
	
	def __init__(self, package):
		self.package = package
		self.run()
	
	def openNetfile(self):
		'''
		Opens a file for writing.
		'''
		self.scanid = self.acquireNumber()
		return NetCDF.NetCDFFile(self.path+self.stm_prefix+self.scanid+self.extension,'w')
	
	def saveWorldSettings(self):
		'''
		Saves parameters from the worldsettings file.
		
		'''
		self.netfile.ANG_PER_VOLT_X = str(ws.ANG_PER_VOLT_X)
		self.netfile.ANG_PER_VOLT_Y = str(ws.ANG_PER_VOLT_Y)
		self.netfile.ANG_PER_VOLT_Z = str(ws.ANG_PER_VOLT_Z)
		self.netfile.VERSION = ws.VERSION
		self.netfile.SCAN_ID = self.scanid
		self.netfile.STM_PREFIX = self.stm_prefix
		
	
	def closeNetfile(self):
		'''
		Closes the netfile.
		'''
		self.netfile.close()
	
	def acquireNumber(self):
		'''
		Retrieves a number from the number file CurrentScan and
		updates the file to the new current scan number. Uses an 8-digit
		save scheme.
		
		'''
		fh = open('CurrentScan','r')
		ret = fh.read()
		fh.close()
		ret = int(ret)
		savestr = '%'+'0.%dd'%ws.SAVE_DIGITS
		while os.path.exists(self.basedir+self.datadir+'/'+self.stm_prefix+str(savestr%ret)+self.extension):
			ret += 1
		num = int(ret)+1
		
		fh = open('CurrentScan','w')
		fh.write(str(num))
		fh.close()
		return str(savestr%ret)
	
	def makeDirs(self):
		'''
		Makes sure that the correct folders exist before saving.
		'''
		self.date = []
		for i in time.localtime(): self.date.append(str(i))
		if len(self.date[1]) == 1: self.date[1] = '0' + self.date[1]
		self.datedir = '/'+self.date[0]+'/'+self.date[1]
		self.path = self.basedir+self.datadir+self.datedir+'/'
		if not os.path.exists(self.path):
			os.makedirs(self.path)
	
	def saveScan(self):
		'''
		This function determines how to save each parameter. Items which are not strings have a dimensions and
		according to the NetCDF standard, must define the dimension before saving. Items which are strings
		are attributes. These are data descriptors.
		'''
		for ele in self.package.__dict__:
			item = self.package.__dict__[ele]
			
			if type(item) == list:
				if type(item[0]) == int:
					#print ele,item[0]
					if (ele == 'SCAN_Y' or ele == 'SCAN_X') and self.package.TYPE in self.imagescans:
						self.netfile.createDimension(ele,item[0])
						self.netfile.createVariable(ele,'i',(ele,))
						self.netfile.variables[ele][:] = N.array(range(1,item[0]+1))
						if ele == 'SCAN_Y':
							self.netfile.PIXELS_Y = str(item[0]) + ' ' +item[1]
						else: self.netfile.PIXELS_X = str(item[0]) + ' ' +item[1]
					else:
						self.netfile.createDimension(ele,1)
						self.netfile.createVariable(ele,'i',(ele,))
						self.netfile.variables[ele][0] = item[0]
					self.netfile.variables[ele].unit = item[1]
				elif type(item[0]) == float:
						self.netfile.createDimension(ele,1)
						self.netfile.createVariable(ele,'f',(ele,))
						self.netfile.variables[ele][0] = item[0]
						self.netfile.variables[ele].unit = item[1]
				elif type(item[0]) == str:
					self.scanAttr(ele,item[0])
			elif type(item) == str:
				self.scanAttr(ele,item)
		# Save the specific data. elif is not used because a scan may save one or more of these
		if self.package.TYPE in self.imagescans: self.saveIMAGEDATA()
		if self.package.TYPE in self.ivscans: self.saveIV()
		if self.package.TYPE in self.izscans: self.saveIZ()
		if self.package.TYPE in self.zvscans: self.saveZV()
		if self.package.TYPE in self.spectroscopy: self.saveMask()
		self.netfile.sync()
	
	def saveIMAGEDATA(self):
		self.netfile.createDimension('DIRECTION',2)
		self.netfile.createDimension('NBIAS_STEPS',len(self.package.IMAGEDATA))
		if hasattr(self.package,'BIAS_STEPS'):
			self.netfile.createVariable('BIAS_STEPS','f',('NBIAS_STEPS',))
			print self.package.BIAS_STEPS
			for i in range(self.package.N_BIAS_STEPS[0]):
				self.netfile.variables['BIAS_STEPS'][i]=self.package.BIAS_STEPS[i][0]
		self.netfile.createVariable('IMAGEDATA','f',('NBIAS_STEPS','DIRECTION','SCAN_Y','SCAN_X'))
		for i in range(len(self.package.IMAGEDATA)):
			for j in range(2):
				self.netfile.variables['IMAGEDATA'][i,j,:,:] = self.package.IMAGEDATA[i][j].tolist()
	
	def saveMask(self):
		self.netfile.createVariable('SPEC_MASK','f',('SCAN_Y','SCAN_X'))
		self.netfile.variables['SPEC_MASK'][:,:] = self.package.SPEC_MASK.tolist()
	
	def saveIV(self):
		coordlen = len(self.package.IV)
		layers = self.package.IV_MULTIZ[0]+1
		self.netfile.createDimension('IV_INDEX',coordlen)
		self.netfile.createDimension('V',self.package.IV_STEPS[0])
		self.netfile.createDimension('N_IV_LAYERS',layers)
		self.netfile.createDimension('IV_DIRECTION',2)
		xiv = []
		yiv = []
		zvals = []
		for ele in self.package.IV:
			yiv.append(int(ele.split(',')[0][1:]))
			xiv.append(int(ele.split(',')[1][:-1]))
		self.netfile.createVariable('xiv','i',('IV_INDEX',))
		self.netfile.createVariable('yiv','i',('IV_INDEX',))
		self.netfile.createVariable('layer','f',('N_IV_LAYERS',))
		self.netfile.createVariable('V','f',('V',))
		#self.netfile.createVariable('iv','f',('V','IV_INDEX','IV_DIRECTION','N_IV_LAYERS',))
		self.netfile.createVariable('iv','f',('N_IV_LAYERS','IV_DIRECTION','IV_INDEX','V',))
		self.netfile.variables['xiv'][:] = xiv
		self.netfile.variables['yiv'][:] = yiv
		self.netfile.variables['V'][:] = n.float32(self.package.IV[str((yiv[0],xiv[0]))][0][0][0])
		self.netfile.variables['V'].unit = 'volts'
		for i in self.package.IV_Z[0:self.package.IV_MULTIZ[0]+1]: zvals.append(i[0])
		self.netfile.variables['layer'][:] = zvals
		self.netfile.variables['layer'].unit = self.package.IV_Z[0][1]
		for k in range(layers):
			for j in range(2):
				for i in range(coordlen):
					if str((yiv[i],xiv[i])) in self.package.IV:
						#self.netfile.variables['iv'][:,i,j,k] = self.package.IV[str((xiv[i],yiv[i]))][k][j][1]
						self.netfile.variables['iv'][k,j,i,:] = n.float32(self.package.IV[str((yiv[i],xiv[i]))][k][j][1])
		
	
	def saveIZ(self):
		niz = len(self.package.IZ)
		izlen = len(self.package.IZ[self.package.IZ.keys()[0]][0][1])
		#x_y = 2
		self.netfile.createDimension('IZ_INDEX',niz)
		self.netfile.createDimension('Z_VALS',self.package.IZ_STEPS[0])
		#self.netfile.createDimension('IZ_XY',x_y)
		self.netfile.createVariable('xiz','i',('IZ_INDEX',))
		self.netfile.createVariable('yiz','i',('IZ_INDEX',))
		self.netfile.createDimension('IZ_DIRECTION',2)
		self.netfile.createVariable('iz','f',('IZ_DIRECTION','IZ_INDEX','Z_VALS',))
		self.netfile.createVariable('Z_VALS','f',('Z_VALS',))
		xiz = []
		yiz = []
		for ele in self.package.IZ:
			yiz.append(int(ele.split(',')[0][1:]))
			xiz.append(int(ele.split(',')[1][:-1]))
		self.netfile.variables['xiz'][:] = xiz
		self.netfile.variables['yiz'][:] = yiz
		self.netfile.variables['Z_VALS'][:] = self.package.IZ[str((yiz[0],xiz[0]))][0][0][0]
		self.netfile.variables['Z_VALS'].unit = 'volts'
		#for i in range(niz):
		for j in range(2):
			for i in range(niz):
				if str((yiz[i],xiz[i])) in self.package.IZ:
					self.netfile.variables['iz'][j,i,:] = self.package.IZ[str((yiz[i],xiz[i]))][0][j][1]
	
	def saveZV(self):
		coordlen = len(self.package.ZV)
		#layers = self.package.ZV_MULTIZ[0]+1
		self.netfile.createDimension('ZV_INDEX',coordlen)
		self.netfile.createDimension('V',self.package.ZV_STEPS[0])
		#self.netfile.createDimension('N_ZV_LAYERS',layers)
		self.netfile.createDimension('ZV_DIRECTION',2)
		xzv = []
		yzv = []
		zvals = []
		for ele in self.package.ZV:
			yzv.append(int(ele.split(',')[0][1:]))
			xzv.append(int(ele.split(',')[1][:-1]))
		self.netfile.createVariable('xzv','i',('ZV_INDEX',))
		self.netfile.createVariable('yzv','i',('ZV_INDEX',))
		#self.netfile.createVariable('layer','f',('N_ZV_LAYERS',))
		self.netfile.createVariable('V','f',('V',))
		#self.netfile.createVariable('zv','f',('V','ZV_INDEX','ZV_DIRECTION','N_ZV_LAYERS',))
		self.netfile.createVariable('zv','f',('ZV_DIRECTION','ZV_INDEX','V',))
		self.netfile.variables['xzv'][:] = xzv
		self.netfile.variables['yzv'][:] = yzv
		self.netfile.variables['V'][:] = self.package.ZV[str((yzv[0],xzv[0]))][0][0][0]
		self.netfile.variables['V'].unit = 'volts'
		for j in range(2):
			for i in range(coordlen):
				if str((yzv[i],xzv[i])) in self.package.ZV:
					#self.netfile.variables['zv'][:,i,j,k] = self.package.ZV[str((xzv[i],yzv[i]))][k][j][1]
					self.netfile.variables['zv'][j,i,:] = self.package.ZV[str((yzv[i],xzv[i]))][0][j][1]
	
	def scanAttr(self,string,item):
		'''
		Saves all attributes according to their name. Attributes are dimensionless,
		single-valued quantities.
		'''
		exec 'self.netfile.%s=item'%string
	
	def run(self):
		'''
		Runs the routine to save all of the data.
		
		'''
		self.makeDirs()
		self.netfile = self.openNetfile()
		self.saveWorldSettings()
		self.saveScan()
		self.closeNetfile()
	